package s3

import (
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/mars/internal/pkg/singleflight"
	"github.com/mars/pkg/maps"
)

const (
	UploadManagerS3ClientKeyPrefix = "uploadManagerClient"
	RawS3ClientKeyPrefix           = "rawS3Client"
)

// NewS3Client create *s3.S3 instance according to host, accessKey and secretKey.
func NewS3Client(host string, ak, sk string) (*s3.S3, error) {
	sess, err := newSession(host, ak, sk)
	if err != nil {
		return nil, err
	}
	return s3.New(sess), nil
}

func NewUploadManager(host string, ak, sk string) (*s3manager.Uploader, error) {
	sess, err := newSession(host, ak, sk)
	if err != nil {
		return nil, err
	}
	return s3manager.NewUploader(sess), nil
}

func newSession(host string, ak, sk string) (*session.Session, error) {
	return session.NewSession(aws.NewConfig().WithEndpoint(host).
		WithMaxRetries(2).
		WithRegion(DefaultRegion).
		WithS3ForcePathStyle(true).
		WithCredentials(credentials.NewStaticCredentials(ak, sk, "")))
}

type S3ClientFactory interface {
	GetS3RawClient() (*s3.S3, error)
	GetUploadManagerClient() (*s3manager.Uploader, error)
}

// S3ClientCache 创建 s3 客户端，因为可能包含多个集群，而且集群会动态变更，所以这里使用懒加载的方式创建集群信息
// 并将集群客户端缓存
type S3ClientCache struct {
	group         singleflight.Group
	s3ClientCache maps.RWMap[string, any] // value 两种类型： *s3.S3 或者 *s3manager.Uploader
}

func (clientSetter *S3ClientCache) getS3Client(key string, server string, ak, sk string) (any, error) {
	client, exists := clientSetter.s3ClientCache.Get(key)
	if exists {
		return client, nil
	}

	return clientSetter.group.Do(key, func() (any, error) {
		var (
			client any
			err    error
		)
		if strings.HasPrefix(key, UploadManagerS3ClientKeyPrefix) {
			client, err = NewUploadManager(server, ak, sk)
			if err != nil {
				return nil, nil
			}
		}
		if strings.HasPrefix(key, RawS3ClientKeyPrefix) {
			client, err = NewS3Client(server, ak, sk)
			if err != nil {
				return nil, err
			}
		}
		if err != nil {
			return nil, err
		}
		clientSetter.s3ClientCache.Put(key, client)
		return client, nil
	})
}

var s3ClientCache = &S3ClientCache{}

type MultiClientFactory struct {
	host string
	ak   string
	sk   string
}

var _ S3ClientFactory = (*MultiClientFactory)(nil)

func NewMultiClientFactory(host string, ak, sk string) *MultiClientFactory {
	return &MultiClientFactory{
		host: host,
		ak:   ak,
		sk:   sk,
	}
}

func (fc *MultiClientFactory) GetUploadManagerClient() (*s3manager.Uploader, error) {
	key := UploadManagerS3ClientKeyPrefix + fc.host + fc.ak
	clientInterface, err := s3ClientCache.getS3Client(key, fc.host, fc.ak, fc.sk)
	if err != nil {
		return nil, err
	}
	return clientInterface.(*s3manager.Uploader), nil
}

func (fc *MultiClientFactory) GetS3RawClient() (*s3.S3, error) {
	key := RawS3ClientKeyPrefix + fc.host + fc.ak
	clientInterface, err := s3ClientCache.getS3Client(key, fc.host, fc.ak, fc.sk)
	if err != nil {
		return nil, err
	}
	return clientInterface.(*s3.S3), nil
}

// SingleClient 只针对一个集群的客户端
type SingleClient struct {
	rawClient           *s3.S3
	uploadManagerClient *s3manager.Uploader
}

func NewSingleClient(host string, ak, sk string) (S3ClientFactory, error) {
	uploadManagerClient, err := NewUploadManager(host, ak, sk)
	if err != nil {
		return nil, nil
	}
	rawClient, err := NewS3Client(host, ak, sk)
	if err != nil {
		return nil, err
	}
	return &SingleClient{rawClient, uploadManagerClient}, nil
}

func (sc *SingleClient) GetUploadManagerClient() (*s3manager.Uploader, error) {
	return sc.uploadManagerClient, nil
}

func (sc *SingleClient) GetS3RawClient() (*s3.S3, error) {
	return sc.rawClient, nil
}
